package net

import (
	"github.com/astaxie/beego/toolbox"
	"github.com/gogf/gf/container/gmap"
	"github.com/rcrowley/go-metrics"
	"github.com/shirou/gopsutil/net"
)

var (
	StatusError error
	Register    metrics.Registry
	History     *gmap.Map
)

func init() {
	History = gmap.New(true)
	Register = metrics.NewRegistry()
	tk := toolbox.NewTask("GetNetDetailedTask", "*/1 * * * * *", GetNetDetailedTask)
	tk2 := toolbox.NewTask("GetFlowTask", "*/1 * * * * *", GetFlowTask)
	toolbox.AddTask("GetNetDetailedTask", tk)
	toolbox.AddTask("GetFlowTask", tk2)
}

func GetNetDetailedTask() error {
	status, err := net.Connections("all")
	if err != nil {
		StatusError = err
		return StatusError
	}
	result := make(map[string]int)
	for _, s := range status {
		if v, ok := result[s.Status]; ok {
			result[s.Status] = v + 1
		} else {
			result[s.Status] = 1
		}
	}
	for kind, count := range result {
		if kind != "" && kind != "NONE" {
			M := metrics.GetOrRegisterGauge(kind, Register)
			M.Update(int64(count))
		}
	}
	return nil
}

func GetFlowTask() error {
	status, err := net.IOCounters(false)
	if err != nil {
		return err
	}
	for _, card := range status {
		bytesRecv := History.GetOrSet("bytes_recv", card.BytesRecv)
		bytesSend := History.GetOrSet("bytes_send", card.BytesSent)
		PacketsRecv := History.GetOrSet("packets_recv", card.PacketsRecv)
		PacketsSent := History.GetOrSet("packets_send", card.PacketsSent)

		metrics.GetOrRegisterGauge("bytes_recv", Register).Update(int64(card.BytesRecv) - int64(bytesRecv.(uint64)))
		metrics.GetOrRegisterGauge("bytes_send", Register).Update(int64(card.BytesSent) - int64(bytesSend.(uint64)))
		metrics.GetOrRegisterGauge("packets_recv", Register).Update(int64(card.PacketsRecv) - int64(PacketsRecv.(uint64)))
		metrics.GetOrRegisterGauge("packets_send", Register).Update(int64(card.PacketsSent) - int64(PacketsSent.(uint64)))
	}
	return nil
}
